

# --- Pinecone 云推理版本 ---
import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone
from collections import Counter

# Pinecone 配置
PINECONE_API_KEY = "pcsk_dqdo1_EP62eXyBSFyv5dQWZiMJZybi6YAkjrT8cYaobXeTkRDvfnbvxY9y8RCdZkajBb1"  
INDEX_NAME = "mnist-index"
K = 11

def preprocess_image(img):
    if isinstance(img, dict):
        if 'composite' in img:
            img = img['composite']
        else:
            raise ValueError("未找到 composite 键，无法处理输入")
    if isinstance(img, np.ndarray):
        if img.shape[-1] == 4:
            rgb, alpha = img[..., :3], img[..., 3]
            bg = np.ones_like(rgb) * 255
            mask = (alpha > 0)[..., None]
            img = np.where(mask, rgb, bg).astype('uint8')
        img = Image.fromarray(img)
    if not isinstance(img, Image.Image):
        raise ValueError(f"无法识别的图像输入类型: {type(img)}")
    img = img.convert('L')
    img = img.resize((8, 8), Image.BILINEAR)
    arr = np.asarray(img).astype(float)
    if arr.max() < 10:
        raise ValueError('画板内容为空，请在画板上写字')
    if arr.mean() > 8:
        arr = 255 - arr
    arr = (arr / 255.0) * 16.0
    arr = arr.reshape(1, -1)
    return arr

def pinecone_predict(img, debug=False):
    try:
        x = preprocess_image(img)
        if x.ndim != 2 or x.shape[1] != 64:
            return (None, f'预处理后形状异常: {x.shape}') if debug else f'Error: preprocessed shape {x.shape}'
        # Pinecone 查询
        pc = Pinecone(api_key=PINECONE_API_KEY)
        index = pc.Index(INDEX_NAME)
        results = index.query(vector=x.flatten().astype(np.float32).tolist(), top_k=K, include_metadata=True)
        if not results['matches']:
            return (None, '未找到匹配') if debug else '未找到匹配'
        labels = [match['metadata']['label'] for match in results['matches']]
        pred = Counter(labels).most_common(1)[0][0]
        if debug:
            stats = {
                'min': float(x.min()),
                'max': float(x.max()),
                'mean': float(x.mean()),
                'raw': np.round(x.reshape(8,8), 2).tolist(),
                'knn_labels': labels
            }
            return pred, f"stats: min={stats['min']:.3f}, max={stats['max']:.3f}, mean={stats['mean']:.3f},\nraw8x8={stats['raw']}\nKNN投票: {labels}"
        return pred
    except Exception as e:
        if debug:
            return None, f'Error: {e}'
        return f'Error: {e}'

def build_gradio_app():
    with gr.Blocks() as demo:
        gr.Markdown('# Pinecone 云 KNN 手写数字识别')
        with gr.Row():
            with gr.Column():
                # 兼容不同 gradio 版本，不加 shape 参数
                sketch = gr.Sketchpad(label='在此画数字 (或上传图片)')
                upload = gr.Image(label='或上传图片', type='pil')
                btn = gr.Button('预测')
            with gr.Column():
                out = gr.Label(label='预测结果')
                debug_out = gr.Textbox(label='Debug 输出（特征统计）')
                debug_chk = gr.Checkbox(label='开启 Debug 输出', value=False)

        def on_click_predict(sketch_img, upload_img, debug):
            img = upload_img if upload_img is not None else sketch_img
            if img is None:
                return (None, '未提供图像') if debug else ('未提供图像', '')
            if debug:
                label, info = pinecone_predict(img, debug=True)
                return label, info
            else:
                result = pinecone_predict(img, debug=False)
                return result, ''

        btn.click(on_click_predict, inputs=[sketch, upload, debug_chk], outputs=[out, debug_out])
    return demo

if __name__ == '__main__':
    try:
        app = build_gradio_app()
        app.launch(share=False)
    except Exception as e:
        print('启动 webapp 时出错:', e)
